from yacs.config import CfgNode as CN

_C = CN()
# -----MODEL
_C.MODEL = CN()

# backbone
_C.MODEL.BACKBONE = CN()
_C.MODEL.BACKBONE.ARCH = 'resnet50'
_C.MODEL.BACKBONE.FREEZE = False
_C.MODEL.BACKBONE.EXTRA = CN(new_allowed=True)

# neck
_C.MODEL.NECK = CN()
_C.MODEL.NECK.ARCH = 'fpn'
_C.MODEL.NECK.IN_CHANNELS = [256, 512, 1024, 2048]
_C.MODEL.NECK.OUT_CHANNELS = [256, 256, 256, 256]
_C.MODEL.NECK.INTERPOLATION_MODE = 'bilinear'
_C.MODEL.NECK.ALIGN_CORNERS = True
_C.MODEL.NECK.USE_BIAS = True

# head
_C.MODEL.HEAD = CN()
_C.MODEL.HEAD.ARCH = 'simple_dilate_head'
_C.MODEL.HEAD.IN_CHANNEL = 1024
_C.MODEL.HEAD.MID_CHANNEL = 256
_C.MODEL.HEAD.NUM_CLASS = 7
_C.MODEL.HEAD.SCALE = 4
_C.MODEL.HEAD.INTERPOLATION_MODE = 'bilinear'
_C.MODEL.HEAD.ALIGN_CORNERS = True
_C.MODEL.HEAD.MIN_SCALE = 0.4  # polygon shrink scale
_C.MODEL.HEAD.USE_BIAS = True

# loss
_C.MODEL.LOSS = CN()
_C.MODEL.LOSS.ARCH = 'dice_loss'

# train hyper-params
_C.MODEL.TRAIN = CN()
_C.MODEL.TRAIN.SUM_STEP = 180000
_C.MODEL.TRAIN.SAVE_PATH = './ckpt/'
_C.MODEL.TRAIN.FINE_TUNE = ''

_C.MODEL.TRAIN.OPTIMIZER = CN()
_C.MODEL.TRAIN.OPTIMIZER.MODE = 'step'
_C.MODEL.TRAIN.OPTIMIZER.LR = 1e-3
_C.MODEL.TRAIN.OPTIMIZER.MOMENTUM = 0.99
_C.MODEL.TRAIN.OPTIMIZER.WEIGHT_DECAY = 5e-4

# test hyper-params
_C.MODEL.TEST = CN()
_C.MODEL.TEST.CKPT_PATH = ''
_C.MODEL.TEST.RES_PATH = './result/'
_C.MODEL.TEST.NULL_JSON_PATH = ''
_C.MODEL.TEST.DECODER_METHOD = 'SIMPLE_DILATE'
_C.MODEL.TEST.OUTPUT_BOX = True

# ------DATASET
_C.DATASET = CN()

# image augmentation
_C.DATASET.AUGMENTATION = CN()
_C.DATASET.AUGMENTATION.RANDOM_RESIZE = CN()
_C.DATASET.AUGMENTATION.RANDOM_RESIZE.RANDOM_SCALES = [0.5, 1.0, 2.0, 3.0]
_C.DATASET.AUGMENTATION.RANDOM_RESIZE.MIN_SIZE = 640

_C.DATASET.AUGMENTATION.RESIZE = CN()
_C.DATASET.AUGMENTATION.RESIZE.LONG_SIZE = 1280

_C.DATASET.AUGMENTATION.RANDOM_ROTATE = CN()
_C.DATASET.AUGMENTATION.RANDOM_ROTATE.BASE_ANGLES = [0] #[-180, -90, 0, 90]
_C.DATASET.AUGMENTATION.RANDOM_ROTATE.MAX_ANGLE = 10

_C.DATASET.AUGMENTATION.RANDOM_CROP = CN()
_C.DATASET.AUGMENTATION.RANDOM_CROP.MIN_SIZE = 640

# train dataset
_C.DATASET.TRAIN = CN()
_C.DATASET.TRAIN.CHECK_ALL_IMGS = False
_C.DATASET.TRAIN.IMG_FORMATS = ['.jpg', '.png', '.jpeg']
_C.DATASET.TRAIN.ROOT_PATH = ''
_C.DATASET.TRAIN.DATASET_NAME_LIST = ['menu_train']

_C.DATASET.TRAIN.DATA_AUGMENTATION = ['random_horizontal_flip', 'random_rotate', 'random_crop']
_C.DATASET.TRAIN.BATCH_SIZE = 16
_C.DATASET.TRAIN.NUM_WORKERS = 16
_C.DATASET.TRAIN.DROP_LAST = True
_C.DATASET.TRAIN.PIN_MEMORY = True

# test dataset
_C.DATASET.TEST = CN()
_C.DATASET.TEST.IMG_FORMATS = ['.jpg', '.png', '.jpeg']
_C.DATASET.TEST.ROOT_PATH = ''
_C.DATASET.TEST.LONG_SIZE = 2240


if __name__ == '__main__':
    print(_C)
